import os
import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.models as models
from matplotlib import pyplot as plt
from hub.yolov8 import yolov8, yolov8_cbam, yolov8_transfromer_head

def draw_feature_map(model, input_tensor, save_path):
    _model = torchvision.models._utils.IntermediateLayerGetter(model, '')


if __name__ == '__main__':
    device = torch.device('cuda')
    model = yolov8.DetectHead(3, 's').to(device)
    pretrained_state_dict = torch.load('weights/SYNTHIA/best_state_dict.pt', map_location=device)
    model.load_state_dict(pretrained_state_dict)
    model.eval()
    
    